#include "rtspcamera.h"

#include "mat.h"

RtspCamera::RtspCamera() {
    ifmt_ctx = nullptr;
    ofmt_ctx = nullptr;
    in_codecpar = nullptr;
    ofmt = nullptr;
    src_frame = av_frame_alloc();
    dst_frame = av_frame_alloc();
}
bool RtspCamera::busyness = false;

RtspCamera::~RtspCamera() = default;


void RtspCamera::log_packet(const AVFormatContext *fmt_ctx, const AVPacket *pkt, const char *tag) {
    AVRational *time_base = &fmt_ctx->streams[pkt->stream_index]->time_base;

    __android_log_print(ANDROID_LOG_DEBUG, "rtsp",
                        "%s: pts:%s pts_time:%s dts:%s dts_time:%s duration:%s duration_time:%s stream_index:%d",
                        tag,
                        av_ts2str(pkt->pts), av_ts2timestr(pkt->pts, time_base),
                        av_ts2str(pkt->dts), av_ts2timestr(pkt->dts, time_base),
                        av_ts2str(pkt->duration), av_ts2timestr(pkt->duration, time_base),
                        pkt->stream_index);
}

int RtspCamera::saveJpg(AVFrame *pFrame, char *out_name) const {
    int ret;
    int width = pFrame->width;
    int height = pFrame->height;
    AVCodecContext *pCodeCtx = nullptr;

    AVFormatContext *pFormatCtx = avformat_alloc_context();
    pFormatCtx->oformat = av_guess_format("mjpeg", nullptr, nullptr);

//    if (avio_open(&pFormatCtx->pb, out_name, AVIO_FLAG_READ_WRITE) < 0) {
//        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Couldn't open output file");
//        return -1;
//    }

    AVStream *pAVStream = avformat_new_stream(pFormatCtx, 0);
    if (pAVStream == nullptr) {
        return -1;
    }

    AVCodecParameters *parameters = pAVStream->codecpar;
    parameters->codec_id = pFormatCtx->oformat->video_codec;
    parameters->codec_type = AVMEDIA_TYPE_VIDEO;
    parameters->format = AV_PIX_FMT_YUVJ420P;
    parameters->width = pFrame->width;
    parameters->height = pFrame->height;

    AVCodec *pCodec = avcodec_find_encoder(pAVStream->codecpar->codec_id);

    if (!pCodec) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Could not find encoder");
        return -1;
    }

    pCodeCtx = avcodec_alloc_context3(pCodec);
    if (!pCodeCtx) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera",
                            "Could not allocate video codec context");
        return -99;
    }

    if ((avcodec_parameters_to_context(pCodeCtx, pAVStream->codecpar)) < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera",
                            "Failed to copy %s codec parameters to decoder context\n",
                            av_get_media_type_string(AVMEDIA_TYPE_VIDEO));
        return -1;
    }

    pCodeCtx->time_base = (AVRational) {1, 25};

    if (avcodec_open2(pCodeCtx, pCodec, nullptr) < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Could not open codec");
        return -1;
    }

//    ret = avformat_write_header(pFormatCtx, nullptr);
//    if (ret < 0) {
//        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "write_header fail");
//        return -1;
//    }

    int y_size = width * height;

    AVPacket pkt;
    av_new_packet(&pkt, y_size * 3);

    ret = avcodec_send_frame(pCodeCtx, pFrame);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Could not avcodec_send_frame");
        return -1;
    }

    ret = avcodec_receive_packet(pCodeCtx, &pkt);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Could not avcodec_receive_packet");
        return -1;
    }

    // ret = av_write_frame(pFormatCtx, &pkt);

    cv::Mat rgb(720, 1280, CV_8UC3);
    __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "pkt.len=%d", pkt.size);
    rgb.data = pkt.data;
    on_image(rgb);

    if (ret < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "Could not av_write_frame");
        return -1;
    }

    av_packet_unref(&pkt);


    // av_write_trailer(pFormatCtx);


    avcodec_close(pCodeCtx);
    avio_close(pFormatCtx->pb);
    avformat_free_context(pFormatCtx);

    return 0;
}


void
RtspCamera::decode(AVCodecContext *dec_ctx, AVFrame *_src_frame, AVFrame *_dst_frame, AVPacket *pkt,
                   const char *filename) const {
    char buf[1024];
    int ret;

    ret = avcodec_send_packet(dec_ctx, pkt);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "Error sending a packet for decoding:%d",
                            ret);
        return;
    }

    while (ret >= 0) {
        ret = avcodec_receive_frame(dec_ctx, _src_frame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
            return;
        else if (ret < 0) {
            __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "Error during decoding");
            return;
        }

        /* the picture is allocated by the decoder. no need to free it */
        snprintf(buf, sizeof(buf), "%s_%d.jpg", filename, dec_ctx->frame_number);
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "saving __src_frame %05d  shape(%d,%d)",
                            dec_ctx->frame_number, src_frame->width, src_frame->height);


        SwsContext *sws_ctx = sws_getContext(1280, 720,
                                             AV_PIX_FMT_YUV420P,
                                             1280, 720,
                                             AV_PIX_FMT_RGB24,
                                             SWS_BICUBIC,
                                             nullptr,
                                             nullptr,
                                             nullptr);

        sws_scale(sws_ctx, _src_frame->data, _src_frame->linesize,
                  0, 720,
                  _dst_frame->data, _dst_frame->linesize);


        cv::Mat rgb(720, 1280, CV_8UC3);
        rgb.data = _dst_frame->data[0];
        __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "src_frame.len=%d", rgb.flags);

        busyness = true;
        on_image(rgb);
        busyness = false;

        // saveJpg(src_frame, buf);
    }
}

bool RtspCamera::open(const char *rtspUrl) {
    __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "open rtsp:%s", rtspUrl);
    running = true;
    av_register_all();
    avformat_network_init();

    const char *outfilename = "/sdcard/img";
    AVPacket *pkt = nullptr;
    int ret;
    pkt = av_packet_alloc();
    if (!pkt) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "pkt error");
        return false;
    }
    AVDictionary *rtsp_option = nullptr;
    av_dict_set(&rtsp_option, "rtsp_transport", "tcp", 0);
    ret = avformat_open_input(&ifmt_ctx, rtspUrl, nullptr, &rtsp_option);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "Could not open '%s' %d", rtspUrl, ret);
        return false;
    }

    if ((ret = avformat_find_stream_info(ifmt_ctx, nullptr)) < 0) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg",
                            "Failed to retrieve input stream information");
        return false;
    }

    av_dump_format(ifmt_ctx, 0, rtspUrl, 0);

    for (int i = 0; i < ifmt_ctx->nb_streams; i++) {
        AVStream *inStream = ifmt_ctx->streams[i];
        in_codecpar = inStream->codecpar;
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "codec_type %d", in_codecpar->codec_type);
        if (in_codecpar->codec_type == AVMEDIA_TYPE_VIDEO) {
            video_trade_index = i;
            break;
        }
    }

    __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "rtsp connect success");

    codec = avcodec_find_decoder(AV_CODEC_ID_H264);
    if (!codec) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Codec not found codec");
        return -4;
    }

    h264 = avcodec_alloc_context3(codec);

    if (!h264) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Could not allocate video codec context");
        return -5;
    }

    /* open it */
    if (avcodec_open2(h264, codec, nullptr) < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Could not open codec");
        return -6;
    }

    src_frame = av_frame_alloc();
    if (!src_frame) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Could not allocate video src_frame");
        return -6;
    }


    int count_bytes_src = avpicture_get_size(AV_PIX_FMT_YUV420P, 1280, 720);
    int count_bytes_dst = avpicture_get_size(AV_PIX_FMT_RGB24, 1280, 720);

    auto *src_buff = (uint8_t *) av_malloc(count_bytes_src);
    auto *dst_buff = (uint8_t *) av_malloc(count_bytes_dst);

    avpicture_fill((AVPicture *) src_frame, src_buff, AV_PIX_FMT_YUV420P, 1280, 720);
    avpicture_fill((AVPicture *) dst_frame, dst_buff, AV_PIX_FMT_RGB24, 1280, 720);


    while (running) {

        AVStream *in_stream, *out_stream;
        av_usleep(10000);
        ret = av_read_frame(ifmt_ctx, pkt);
        if (ret == AVERROR_EOF) {
            __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "NO FRAME:%d", ret);
            break;
        }
        in_stream = ifmt_ctx->streams[pkt->stream_index];
        if (pkt->stream_index != video_trade_index) {
            av_packet_unref(pkt);
            continue;
        }

        if (recording && ofmt_ctx) {
            if (!busyness){
                decode(h264, src_frame, dst_frame, pkt, outfilename);
            }
//            if (!frameKeyStart) {
//                frameKeyStart = pkt->flags == AV_PKT_FLAG_KEY;
//            }
//            if (frameKeyStart) {
//                pkt->stream_index = video_trade_index;
//                out_stream = ofmt_ctx->streams[pkt->stream_index];
//                log_packet(ifmt_ctx, pkt, "in");
//
//                /* copy packet */
//                av_packet_rescale_ts(pkt, in_stream->time_base, out_stream->time_base);
//                pkt->pos = -1;
//                log_packet(ofmt_ctx, pkt, "out");
//
//                ret = av_interleaved_write_frame(ofmt_ctx, pkt);
//
//                /* pkt is now blank (av_interleaved_write_frame() takes ownership of
//                 * its contents and resets pkt), so that no unreferencing is necessary.
//                 * This would be different if one used av_write_frame(). */
//                if (ret < 0) {
//                    __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Error muxing packet:%d", ret);
//                }
//            }
        }
    }
    __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "摄像头即将释放");
    av_packet_free(&pkt);
    avformat_close_input(&ifmt_ctx);
    if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE)) {
        avio_closep(&ofmt_ctx->pb);
    }
    if (ofmt_ctx) avformat_free_context(ofmt_ctx);
    if (ret < 0 && ret != AVERROR_EOF) {
        return false;
    }
    return true;
}

void RtspCamera::close() {
    running = false;
}

void RtspCamera::start(const char *out_filename) {
    if (ofmt_ctx || video_trade_index < 0) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "NOT READY");
    }
    int ret;
    __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "创建输出格式");
    ret = avformat_alloc_output_context2(&ofmt_ctx, nullptr, nullptr, out_filename);
    if (!ofmt_ctx) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Could not create output context %d %d",
                            video_trade_index, ret);
        return;
    }

    ofmt = ofmt_ctx->oformat;
    av_dump_format(ofmt_ctx, 0, out_filename, 1);
    if (!(ofmt->flags & AVFMT_NOFILE)) {
        ret = avio_open(&ofmt_ctx->pb, out_filename, AVIO_FLAG_WRITE);
        if (ret < 0) {
            __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Could not open output file '%s'",
                                out_filename);
            return;
        }
    }


    AVStream *outStream;
    outStream = avformat_new_stream(ofmt_ctx, nullptr);
    if (!outStream) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Failed allocating output stream");
        return;
    }

    ret = avcodec_parameters_copy(outStream->codecpar, in_codecpar);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "Failed to copy codec parameters");
        return;
    }
    outStream->codecpar->codec_tag = 0;

    ret = avformat_write_header(ofmt_ctx, nullptr);
    if (ret < 0) {
        __android_log_print(ANDROID_LOG_ERROR, "ffmpeg",
                            "Error occurred when opening output file:%d", ret);
        return;
    }
    frameKeyStart = false;
    recording = true;
    __android_log_print(ANDROID_LOG_DEBUG, "ffmpeg", "文件路径 >>> %s", out_filename);
}

void RtspCamera::stop() {
    recording = false;
    av_write_trailer(ofmt_ctx);
    if (ofmt_ctx && !(ofmt->flags & AVFMT_NOFILE))
        avio_closep(&ofmt_ctx->pb);
    avformat_free_context(ofmt_ctx);
    ofmt_ctx = nullptr;
    __android_log_print(ANDROID_LOG_ERROR, "ffmpeg", "停止录制");
}

void RtspCamera::set_status(int detect_status) {

}

void RtspCamera::on_image(const cv::Mat &rgb) const {
    __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "RtspCamera on_image");


}

void RtspCameraWindow::on_image(const cv::Mat &rgb) const {
    __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "RtspCameraWindow on_image");


    on_image_render(rgb);

    int render_h = 720;
    int render_w = 1280;
    // rotate to native window orientation
    cv::Mat rgb_render(render_h, render_w, CV_8UC3);
    ncnn::kanna_rotate_c3(rgb.data, 1280, 720, rgb_render.data, render_w, render_h, 1);

    ANativeWindow_setBuffersGeometry(win, render_w, render_h,
                                     AHARDWAREBUFFER_FORMAT_R8G8B8A8_UNORM);

    ANativeWindow_Buffer buf;
    ANativeWindow_lock(win, &buf, NULL);

    // scale to target size
    if (buf.format == AHARDWAREBUFFER_FORMAT_R8G8B8A8_UNORM ||
        buf.format == AHARDWAREBUFFER_FORMAT_R8G8B8X8_UNORM) {
        for (int y = 0; y < render_h; y++) {
            const unsigned char *ptr = rgb_render.ptr<const unsigned char>(y);
            unsigned char *outptr = (unsigned char *) buf.bits + buf.stride * 4 * y;

            int x = 0;
#if __ARM_NEON
            for (; x + 7 < render_w; x += 8) {
                uint8x8x3_t _rgb = vld3_u8(ptr);
                uint8x8x4_t _rgba;
                _rgba.val[0] = _rgb.val[0];
                _rgba.val[1] = _rgb.val[1];
                _rgba.val[2] = _rgb.val[2];
                _rgba.val[3] = vdup_n_u8(255);
                vst4_u8(outptr, _rgba);

                ptr += 24;
                outptr += 32;
            }
#endif // __ARM_NEON
            for (; x < render_w; x++) {
                outptr[0] = ptr[0];
                outptr[1] = ptr[1];
                outptr[2] = ptr[2];
                outptr[3] = 255;

                ptr += 3;
                outptr += 4;
            }
        }
    }

    ANativeWindow_unlockAndPost(win);
}

void RtspCameraWindow::on_image_render(const cv::Mat &rgb) const {
    __android_log_print(ANDROID_LOG_DEBUG, "rtspcamera", "RtspCameraWindow on_image_render");
}

RtspCameraWindow::RtspCameraWindow() {
    win = 0;
}

RtspCameraWindow::~RtspCameraWindow() {
    if (win) {
        ANativeWindow_release(win);
    }
}

void RtspCameraWindow::set_window(ANativeWindow *_win) {
    if (win) {
        ANativeWindow_release(win);
    }
    win = _win;
    ANativeWindow_acquire(win);
}

